#!/usr/bin/python

"""
This program draws gene clusters to make publication quality figures. It takes in
Genbank file, COG category file, and expects start and stop coordinates of region
to inspect.

Usage: dissertation_DrawGenes.py seq.gbk cogs.t.list 10000 20000
Examples:
dissertation_DrawGenesArrows.py ../../../annotation/GKIL.v6.gbf cogs.t.list 10000 20000

Note: Resolution is best if the segment in view is less than 50000bp.

Author: Jimmy Saw
Date of last update: 04-23-2012

"""

import sys
import re
import matplotlib.pyplot as plt
import pylab
import matplotlib
from matplotlib import mpl
from matplotlib.patches import Rectangle
from matplotlib.transforms import Bbox
from Bio import SeqIO
from Bio.SeqUtils import GC
import matplotlib.patches as mpatch

#Regex and other stuffs
cogcat = re.compile('\[(.*)\]\t(\w+)\t.*')

cogdict = {
    'J' : '#2B60DE',
    'A' : '#F6358A',
    'K' : '#B048B5',
    'L' : '#8E35EF',
    'B' : '#D16587',
    'D' : '#C38EC7',
    'Y' : '#52F3FF',
    'V' : '#3EA99F',
    'T' : '#254117',
    'M' : '#41A317',
    'N' : '#00FF00',
    'Z' : '#FFFF00',
    'W' : '#FDD017',
    'U' : '#F88017',
    'O' : '#F660AB',
    'C' : '#FF0000',
    'G' : '#FAAFBA',
    'E' : '#7F5A58',
    'F' : '#C8B560',
    'H' : '#D2B9D3',
    'I' : '#C12869',
    'P' : '#57E964',
    'Q' : '#BCE954',
    'R' : '#F87431',
    'S' : '#ADA96E',
    'X' : '#000000' #No COG
}

##Genome one Genbank file
g1seq = SeqIO.read(sys.argv[1], "gb")
g1length = len(g1seq.seq)
g1featdict = {}
for feat in g1seq.features:
    if feat.type == 'CDS':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'tRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'rRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat

cogcatfile = open(sys.argv[2], "rU")
cfl = cogcatfile.readlines()

cogcatdict = {}

for line in cfl:
    tmp = line.strip()
    if cogcat.match(tmp):
        pattern = cogcat.match(tmp)
        cogcatdict[pattern.group(2)] = pattern.group(1)[0]

cogcatfile.close()

spanx1 = int(sys.argv[3])
spanx2 = int(sys.argv[4])

glist = []

genome1x = [spanx1, spanx2]
genome1y = [2, 2]
genome2x = [spanx1, spanx2]
genome2y = [10, 10]

##Start plotting
fig = plt.figure(1, figsize=(16,4))
#ax1 = fig.add_subplot(211) #makes the subplot and squeezes the figure to half panel
ax1 = fig.add_subplot(111) #makes the full figure plot. larger.
ax1.plot(genome1x, genome1y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
ax1.plot(genome2x, genome2y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)

#ax1.axis([0, g1length, 0, 14])
ax1.axis([spanx1, spanx2, 0, 14])

for k, v in g1featdict.iteritems():
    g1feat = v
    g1start = g1feat.location._start.position
    g1stop = g1feat.location._end.position
    g1size = g1stop - g1start + 1
    g1mid = g1start + ((g1stop - g1start) / 2.0)
    g1desc = g1feat.qualifiers['product'][0]
    g1gene = ""
    if g1feat.qualifiers.has_key('gene'):
        g1gene = g1feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g1gene = g1feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = "#000000" #base color
    if g1feat.qualifiers.has_key('note'):
        cog = g1feat.qualifiers['note'][0]
        if cog in cogcatdict:
            cogcolor = cogdict[cogcatdict[cog]]
    if g1feat.type == 'tRNA':
        cogcolor = '#800000'
    if g1feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g1gene = g1desc
    if g1feat.strand == -1:
        #rect = Rectangle((g1start, 6.0), g1size, 0.5, fc=cogcolor, ec=cogcolor, alpha=0.5)
        #plt.gca().add_patch(rect)
        x1a = g1start + (g1size * 0.1) #add 10% to make a polygon 
        x1b = g1start
        x2b = g1stop
        x2a = g1stop
        x = [x1a, x1b, x2b, x2a]
        y = [6, 6.5, 6.5, 6]
        if spanx1 < g1start < spanx2 and spanx1 < g1stop < spanx2:
            ax1.fill(x, y, fc=cogcolor, ec=cogcolor, alpha=0.5)
            ax1.text(g1mid, 6.5, g1gene, fontsize=8, color='black', rotation=45)
        #ax1.plot(g1start, 6.0, 'r<', mec='red')
    else:
        #rect = Rectangle((g1start, 7.0), g1size, 0.5, fc=cogcolor, ec=cogcolor, alpha=0.5)
        #plt.gca().add_patch(rect)
        x1a = g1start
        x1b = g1start
        x2b = g1stop - (g1size * 0.1) #substract 10% to make a polygon
        x2a = g1stop
        x = [x1a, x1b, x2b, x2a]
        y = [7, 7.5, 7.5, 7]
        if spanx1 < g1start < spanx2 and spanx1 < g1stop < spanx2:
            ax1.fill(x, y, fc=cogcolor, ec=cogcolor, alpha=0.5)
            ax1.text(g1mid, 7.5, g1gene, fontsize=8, color='black', rotation=45)
        #ax1.plot(g1stop, 7.0, 'r>', mec='red')

"""
#This segment is an example of how to draw polygon connecting genomic regions
a = [140000, 140000, 141000-(1000 * 0.1), 141000] # + strand
b = [5, 5.5, 5.5, 5]
c = [141000+(1000*0.1), 141000, 142000, 142000] # - strand
d = [4.5, 5, 5, 4.5]

ax1.fill(a, b, fc='red', ec='red', alpha=0.2)
ax1.fill(c, d, fc='blue', ec='blue', alpha=0.2)
"""


"""
ax1.text(1710000, 3, 'Test',
          ha="center",
          size=30,
          bbox=dict(boxstyle='rarrow', fc="w", ec="k"))
"""


frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
ax1.grid(False)


plt.show()


